{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DDJwQPZcupab"
   },
   "source": [
    "# EECS 498-007/598-005 Assignment 4-2: RNN, LSTM, and Attention for Image Captioning\n",
    "\n",
    "Before we start, please put your name and UMID in following format\n",
    "\n",
    ": Firstname LASTNAME, #00000000   //   e.g.) Justin JOHNSON, #12345678"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2KMxqLt1h2kx"
   },
   "source": [
    "**Your Answer:**   \n",
    "Hello WORLD, #XXXXXXXX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Euek3FWn6bhA",
    "tags": [
     "pdf-title"
    ]
   },
   "source": [
    "# Image Captioning with RNNs\n",
    "In this exercise you will implement a vanilla recurrent neural networks and use them it to train a model that can generate novel captions for images."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MzqbYcKdz6ew"
   },
   "source": [
    "# Setup Code\n",
    "Before getting started we need to run some boilerplate code to set up our environment. You'll need to rerun this setup code each time you start the notebook.\n",
    "\n",
    "First, run this cell load the [autoreload](https://ipython.readthedocs.io/en/stable/config/extensions/autoreload.html?highlight=autoreload) extension. This allows us to edit `.py` source files, and re-import them into the notebook for a seamless editing and debugging experience."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "28O_qwFfdQpr"
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "p3H9pcnudWlg"
   },
   "source": [
    "### Google Colab Setup\n",
    "\n",
    "Next we need to run a few commands to set up our environment on Google Colab. If you are running this notebook on a local machine you can skip this section.\n",
    "\n",
    "Run the following cell to mount your Google Drive. Follow the link, sign in to your Google account (the same account you used to store this notebook!) and copy the authorization code into the text box that appears below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Yv8Z8EiudX25"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Sek0GtVOdlKT"
   },
   "source": [
    "Now recall the path in your Google Drive where you uploaded this notebook, fill it in below. If everything is working correctly then running the folowing cell should print the filenames from the assignment:\n",
    "\n",
    "```\n",
    "['eecs598', 'network_visualization.py', 'style_transfer.py',  'network_visualization.ipynb', 'a4_helper.py', 'pytorch_autograd_and_nn.py', 'pytorch_autograd_and_nn.ipynb', 'style_transfer.ipynb', 'rnn_lstm_attention_captioning.ipynb',  'rnn_lstm_attention_captioning.py']\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A9t0-bGZdnr8"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# TODO: Fill in the Google Drive path where you uploaded the assignment\n",
    "# Example: If you create a 2020FA folder and put all the files under A1 folder, then '2020FA/A1'\n",
    "GOOGLE_DRIVE_PATH_AFTER_MYDRIVE = None\n",
    "GOOGLE_DRIVE_PATH = os.path.join('drive', 'My Drive', GOOGLE_DRIVE_PATH_AFTER_MYDRIVE)\n",
    "print(os.listdir(GOOGLE_DRIVE_PATH))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "S5LWJPBtdrpZ"
   },
   "source": [
    "Once you have successfully mounted your Google Drive and located the path to this assignment, run th following cell to allow us to import from the `.py` files of this assignment. If it works correctly, it should print the message:\n",
    "\n",
    "```\n",
    "Hello from rnn_lstm_attention_captioning.py!\n",
    "```\n",
    "\n",
    "as well as the last edit time for the file `rnn_lstm_attention_captioning.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jFZKi0podzhO"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(GOOGLE_DRIVE_PATH)\n",
    "\n",
    "import time, os\n",
    "os.environ[\"TZ\"] = \"US/Eastern\"\n",
    "time.tzset() \n",
    "\n",
    "from rnn_lstm_attention_captioning import *\n",
    "from a4_helper import *\n",
    "hello()\n",
    "\n",
    "py_path = os.path.join(GOOGLE_DRIVE_PATH, 'rnn_lstm_attention_captioning.py')\n",
    "py_edit_time = time.ctime(os.path.getmtime(py_path))\n",
    "print('rnn_lstm_attention_captioning.py last edited on %s' % py_edit_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l8fTwRpXfwyM"
   },
   "source": [
    "### Load Packages\n",
    "\n",
    "Run some setup code for this notebook: Import some useful packages and increase the default figure size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "q53DlMXboP-T"
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn.parameter import Parameter\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from eecs598.utils import reset_seed, tensor_to_image, decode_captions, attention_visualizer\n",
    "from eecs598.grad import rel_error, compute_numeric_gradient\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "\n",
    "# for plotting\n",
    "%matplotlib inline\n",
    "plt.rcParams['figure.figsize'] = (10.0, 8.0) # set default size of plots\n",
    "plt.rcParams['image.interpolation'] = 'nearest'\n",
    "plt.rcParams['image.cmap'] = 'gray'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OvUDZWGU3VLV"
   },
   "source": [
    "We will use GPUs to accelerate our computation in this notebook. Run the following to make sure GPUs are enabled:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RrAX9FOLpr9k"
   },
   "outputs": [],
   "source": [
    "if torch.cuda.is_available:\n",
    "  print('Good to go!')\n",
    "else:\n",
    "  print('Please set GPU via Edit -> Notebook Settings.')\n",
    "\n",
    "# data type and device for torch.tensor\n",
    "to_float = {'dtype': torch.float, 'device': 'cpu'}\n",
    "to_float_cuda = {'dtype': torch.float, 'device': 'cuda'}\n",
    "to_double = {'dtype': torch.double, 'device': 'cpu'}\n",
    "to_double_cuda = {'dtype': torch.double, 'device': 'cuda'}\n",
    "to_long = {'dtype': torch.long, 'device': 'cpu'}\n",
    "to_long_cuda = {'dtype': torch.long, 'device': 'cuda'}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WCPZwvOd6bhF",
    "tags": [
     "pdf-ignore"
    ]
   },
   "source": [
    "# Microsoft COCO\n",
    "For this exercise we will use the 2014 release of the [Microsoft COCO dataset](http://mscoco.org/) which has become the standard testbed for image captioning. The dataset consists of 80,000 training images and 40,000 validation images, each annotated with 5 captions written by workers on Amazon Mechanical Turk.\n",
    "\n",
    "We have preprocessed the data for you already and saved them into a serialized data file. It contains 10,000 image-caption pairs for training and 500 for testing. The images have been downsampled to 112x112 for computation efficiency and captions are tokenized and numericalized, clamped to 15 words. You can download the file named `coco.pt` (378MB) with the link below and run some useful stats.\n",
    "\n",
    "You will later use MobileNet v2 to extract features for the images. A few notes on the caption preprocessing:\n",
    "\n",
    "Dealing with strings is inefficient, so we will work with an encoded version of the captions. Each word is assigned an integer ID, allowing us to represent a caption by a sequence of integers. The mapping between integer IDs and words is saved in an entry named `vocab` (both `idx_to_token` and `token_to_idx`), and we use the function `decode_captions` from `a4_helper.py` to convert tensors of integer IDs back into strings.\n",
    "\n",
    "There are a couple special tokens that we add to the vocabulary. We prepend a special `<START>` token and append an `<END>` token to the beginning and end of each caption respectively. Rare words are replaced with a special `<UNK>` token (for \"unknown\"). In addition, since we want to train with minibatches containing captions of different lengths, we pad short captions with a special `<NULL>` token after the `<END>` token and don't compute loss or gradient for `<NULL>` tokens. Since they are a bit of a pain, we have taken care of all implementation details around special tokens for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IMok4gFXqjre"
   },
   "outputs": [],
   "source": [
    "# Download and load serialized COCO data from coco.pt\n",
    "# It contains a dictionary of\n",
    "# \"train_images\" - resized training images (112x112)\n",
    "# \"val_images\" - resized validation images (112x112)\n",
    "# \"train_captions\" - tokenized and numericalized training captions\n",
    "# \"val_captions\" - tokenized and numericalized validation captions\n",
    "# \"vocab\" - caption vocabulary, including \"idx_to_token\" and \"token_to_idx\"\n",
    "\n",
    "if os.path.isfile('./datasets/coco.pt'):\n",
    "  print('COCO data exist')\n",
    "else:\n",
    "  print('downloading COCO dataset')\n",
    "  !wget http://web.eecs.umich.edu/~justincj/teaching/eecs498/coco.pt -P ./datasets/\n",
    "\n",
    "# load COCO data from coco.pt, loaf_COCO is implemented in a4_helper.py\n",
    "data_dict = load_COCO(path = './datasets/coco.pt')\n",
    "\n",
    "num_train = data_dict['train_images'].size(0)\n",
    "num_val = data_dict['val_images'].size(0)\n",
    "\n",
    "# declare variables for special tokens\n",
    "NULL_index = data_dict['vocab']['token_to_idx']['<NULL>']\n",
    "START_index = data_dict['vocab']['token_to_idx']['<START>']\n",
    "END_index = data_dict['vocab']['token_to_idx']['<END>']\n",
    "UNK_index = data_dict['vocab']['token_to_idx']['<UNK>']\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "80RW_nSH6bhH"
   },
   "source": [
    "## Look at the data\n",
    "It is always a good idea to look at examples from the dataset before working with it.\n",
    "\n",
    "Run the following to sample a small minibatch of training data and show the images and their captions. Running it multiple times and looking at the results helps you to get a sense of the dataset.\n",
    "\n",
    "Note that we decode the captions using the `decode_captions` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "l-oiW9Ut6bhH"
   },
   "outputs": [],
   "source": [
    "# Sample a minibatch and show the reshaped 112x112 images and captions\n",
    "batch_size = 3\n",
    "\n",
    "sample_idx = torch.randint(0, num_train, (batch_size,))\n",
    "sample_images = data_dict['train_images'][sample_idx]\n",
    "sample_captions = data_dict['train_captions'][sample_idx]\n",
    "for i in range(batch_size):\n",
    "  plt.imshow(sample_images[i].permute(1, 2, 0))\n",
    "  plt.axis('off')\n",
    "  caption_str = decode_captions(sample_captions[i], data_dict['vocab']['idx_to_token'])\n",
    "  plt.title(caption_str)\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "b2SQMNIH6bhJ"
   },
   "source": [
    "# Recurrent Neural Networks\n",
    "As discussed in lecture, we will use Recurrent Neural Network (RNN) language models for image captioning. We will cover the vanilla RNN model first and later LSTM and attention-based language models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6XHZMI356bhJ"
   },
   "source": [
    "## Vanilla RNN: step forward\n",
    "First implement the `rnn_step_forward` for a single timestep of a vanilla recurrent neural network. Run the following to check your implementation. You should see errors on the order of `1e-8` or less."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "c3oU8JJj6bhK"
   },
   "outputs": [],
   "source": [
    "N, D, H = 3, 10, 4\n",
    "\n",
    "x = torch.linspace(-0.4, 0.7, steps=N*D, **to_double_cuda).reshape(N, D)\n",
    "prev_h = torch.linspace(-0.2, 0.5, steps=N*H, **to_double_cuda).reshape(N, H)\n",
    "Wx = torch.linspace(-0.1, 0.9, steps=D*H, **to_double_cuda).reshape(D, H)\n",
    "Wh = torch.linspace(-0.3, 0.7, steps=H*H, **to_double_cuda).reshape(H, H)\n",
    "b = torch.linspace(-0.2, 0.4, steps=H, **to_double_cuda)\n",
    "\n",
    "# YOUR_TURN: Impelement rnn_step_forward \n",
    "next_h, _ = rnn_step_forward(x, prev_h, Wx, Wh, b)\n",
    "expected_next_h = torch.tensor([\n",
    "  [-0.58172089, -0.50182032, -0.41232771, -0.31410098],\n",
    "  [ 0.66854692,  0.79562378,  0.87755553,  0.92795967],\n",
    "  [ 0.97934501,  0.99144213,  0.99646691,  0.99854353]], **to_double_cuda)\n",
    "\n",
    "print('next_h error: ', rel_error(expected_next_h, next_h))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tid-ljPA6bhL"
   },
   "source": [
    "## Vanilla RNN: step backward\n",
    "Then implement the `rnn_step_backward` for a single timestep of a vanilla recurrent neural network. Run the following to numerically gradient check your implementation. You should see errors on the order of `1e-8` or less."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KPyfJofC6bhM"
   },
   "outputs": [],
   "source": [
    "reset_seed(0)\n",
    "N, D, H = 4, 5, 6\n",
    "x = torch.randn(N, D, **to_double_cuda)\n",
    "h = torch.randn(N, H, **to_double_cuda)\n",
    "Wx = torch.randn(D, H, **to_double_cuda)\n",
    "Wh = torch.randn(H, H, **to_double_cuda)\n",
    "b = torch.randn(H, **to_double_cuda)\n",
    "\n",
    "out, cache = rnn_step_forward(x, h, Wx, Wh, b)\n",
    "\n",
    "dnext_h = torch.randn(*out.shape, **to_double_cuda)\n",
    "\n",
    "fx = lambda x: rnn_step_forward(x, h, Wx, Wh, b)[0]\n",
    "fh = lambda h: rnn_step_forward(x, h, Wx, Wh, b)[0]\n",
    "fWx = lambda Wx: rnn_step_forward(x, h, Wx, Wh, b)[0]\n",
    "fWh = lambda Wh: rnn_step_forward(x, h, Wx, Wh, b)[0]\n",
    "fb = lambda b: rnn_step_forward(x, h, Wx, Wh, b)[0]\n",
    "\n",
    "dx_num = compute_numeric_gradient(fx, x, dnext_h)\n",
    "dprev_h_num = compute_numeric_gradient(fh, h, dnext_h)\n",
    "dWx_num = compute_numeric_gradient(fWx, Wx, dnext_h)\n",
    "dWh_num = compute_numeric_gradient(fWh, Wh, dnext_h)\n",
    "db_num = compute_numeric_gradient(fb, b, dnext_h)\n",
    "\n",
    "# YOUR_TURN: Impelement rnn_step_backward\n",
    "dx, dprev_h, dWx, dWh, db = rnn_step_backward(dnext_h, cache)\n",
    "\n",
    "print('dx error: ', rel_error(dx_num, dx))\n",
    "print('dprev_h error: ', rel_error(dprev_h_num, dprev_h))\n",
    "print('dWx error: ', rel_error(dWx_num, dWx))\n",
    "print('dWh error: ', rel_error(dWh_num, dWh))\n",
    "print('db error: ', rel_error(db_num, db))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vjZjH5JW6bhN"
   },
   "source": [
    "## Vanilla RNN: forward\n",
    "Now that you have implemented the forward and backward passes for a single timestep of a vanilla RNN, you will combine these pieces to implement a RNN that processes an entire sequence of data. First implement `rnn_forward` by making calls to the `rnn_step_forward` function that you defined earlier.\n",
    "\n",
    "Run the following to check your implementation. You should see errors on the order of `1e-6` or less.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_GQWEn3Z6bhO"
   },
   "outputs": [],
   "source": [
    "N, T, D, H = 2, 3, 4, 5\n",
    "\n",
    "x = torch.linspace(-0.1, 0.3, steps=N*T*D, **to_double_cuda).reshape(N, T, D)\n",
    "h0 = torch.linspace(-0.3, 0.1, steps=N*H, **to_double_cuda).reshape(N, H)\n",
    "Wx = torch.linspace(-0.2, 0.4, steps=D*H, **to_double_cuda).reshape(D, H)\n",
    "Wh = torch.linspace(-0.4, 0.1, steps=H*H, **to_double_cuda).reshape(H, H)\n",
    "b = torch.linspace(-0.7, 0.1, steps=H, **to_double_cuda)\n",
    "\n",
    "# YOUR_TURN: Impelement rnn_forward\n",
    "h, _ = rnn_forward(x, h0, Wx, Wh, b)\n",
    "expected_h = torch.tensor([\n",
    "  [\n",
    "    [-0.42070749, -0.27279261, -0.11074945,  0.05740409,  0.22236251],\n",
    "    [-0.39525808, -0.22554661, -0.0409454,   0.14649412,  0.32397316],\n",
    "    [-0.42305111, -0.24223728, -0.04287027,  0.15997045,  0.35014525],\n",
    "  ],\n",
    "  [\n",
    "    [-0.55857474, -0.39065825, -0.19198182,  0.02378408,  0.23735671],\n",
    "    [-0.27150199, -0.07088804,  0.13562939,  0.33099728,  0.50158768],\n",
    "    [-0.51014825, -0.30524429, -0.06755202,  0.17806392,  0.40333043]]], **to_double_cuda)\n",
    "print('h error: ', rel_error(expected_h, h))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "P570PTsw6bhP"
   },
   "source": [
    "## Vanilla RNN: backward\n",
    "Implement the `rnn_backward` for a vanilla RNN. This should run back-propagation over the entire sequence, making calls to the `rnn_step_backward` function that you defined earlier. You should see errors on the order of `1e-6` or less.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ny25RusA6bhQ"
   },
   "outputs": [],
   "source": [
    "reset_seed(0)\n",
    "\n",
    "N, D, T, H = 2, 3, 10, 5\n",
    "\n",
    "x = torch.randn(N, T, D, **to_double_cuda)\n",
    "h0 = torch.randn(N, H, **to_double_cuda)\n",
    "Wx = torch.randn(D, H, **to_double_cuda)\n",
    "Wh = torch.randn(H, H, **to_double_cuda)\n",
    "b = torch.randn(H, **to_double_cuda)\n",
    "\n",
    "out, cache = rnn_forward(x, h0, Wx, Wh, b)\n",
    "\n",
    "dout = torch.randn(*out.shape, **to_double_cuda)\n",
    "\n",
    "# YOUR_TURN: Impelement rnn_backward\n",
    "dx, dh0, dWx, dWh, db = rnn_backward(dout, cache)\n",
    "\n",
    "fx = lambda x: rnn_forward(x, h0, Wx, Wh, b)[0]\n",
    "fh0 = lambda h0: rnn_forward(x, h0, Wx, Wh, b)[0]\n",
    "fWx = lambda Wx: rnn_forward(x, h0, Wx, Wh, b)[0]\n",
    "fWh = lambda Wh: rnn_forward(x, h0, Wx, Wh, b)[0]\n",
    "fb = lambda b: rnn_forward(x, h0, Wx, Wh, b)[0]\n",
    "\n",
    "dx_num = compute_numeric_gradient(fx, x, dout)\n",
    "dh0_num = compute_numeric_gradient(fh0, h0, dout)\n",
    "dWx_num = compute_numeric_gradient(fWx, Wx, dout)\n",
    "dWh_num = compute_numeric_gradient(fWh, Wh, dout)\n",
    "db_num = compute_numeric_gradient(fb, b, dout)\n",
    "\n",
    "print('dx error: ', rel_error(dx_num, dx))\n",
    "print('dh0 error: ', rel_error(dh0_num, dh0))\n",
    "print('dWx error: ', rel_error(dWx_num, dWx))\n",
    "print('dWh error: ', rel_error(dWh_num, dWh))\n",
    "print('db error: ', rel_error(db_num, db))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oEDUmZOkU_LO"
   },
   "source": [
    "## Vanilla RNN: backward with autograd\n",
    "Now it's time to introduce the lifesaver PyTorch Automatic Differantiation package - `torch.autograd`!\n",
    "\n",
    "`torch.autograd` provides classes and functions implementing **automatic differentiation** of arbitrary scalar valued functions. It requires minimal changes to the existing code - if you pass tensors with `requires_grad=True` to the forward function you wrote earlier, you can just call `.backward(gradient=grad)` on the output to compute gradients on the input and weights.\n",
    "\n",
    "Now we can compare the manual backward pass with the autograd backward pass. **Read through the following.**  You should get a relative error less than `1e-12`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5AMXoqNOVRa_"
   },
   "outputs": [],
   "source": [
    "reset_seed(0)\n",
    "\n",
    "N, D, T, H = 2, 3, 10, 5\n",
    "\n",
    "# set requires_grad=True\n",
    "x = torch.randn(N, T, D, **to_double_cuda, requires_grad=True)\n",
    "h0 = torch.randn(N, H, **to_double_cuda, requires_grad=True)\n",
    "Wx = torch.randn(D, H, **to_double_cuda, requires_grad=True)\n",
    "Wh = torch.randn(H, H, **to_double_cuda, requires_grad=True)\n",
    "b = torch.randn(H, **to_double_cuda, requires_grad=True)\n",
    "\n",
    "out, cache = rnn_forward(x, h0, Wx, Wh, b)\n",
    "\n",
    "dout = torch.randn(*out.shape, **to_double_cuda)\n",
    "\n",
    "# manual backward\n",
    "with torch.no_grad():\n",
    "  dx, dh0, dWx, dWh, db = rnn_backward(dout, cache)\n",
    "\n",
    "# backward with autograd\n",
    "out.backward(dout) # the magic happens here!\n",
    "dx_auto, dh0_auto, dWx_auto, dWh_auto, db_auto = \\\n",
    "  x.grad, h0.grad, Wx.grad, Wh.grad, b.grad\n",
    "\n",
    "print('dx error: ', rel_error(dx_auto, dx))\n",
    "print('dh0 error: ', rel_error(dh0_auto, dh0))\n",
    "print('dWx error: ', rel_error(dWx_auto, dWx))\n",
    "print('dWh error: ', rel_error(dWh_auto, dWh))\n",
    "print('db error: ', rel_error(db_auto, db))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zgmxOjX0prA4"
   },
   "source": [
    "## RNN Module\n",
    "We can now wrap the vanilla RNN we wrote into an `nn.Module`. `nn.Module` is a base class for all neural network modules, more details regarding its attributes, functions, and methods could be found [here](https://pytorch.org/docs/stable/nn.html?highlight=module#torch.nn.Module).\n",
    "\n",
    "Here we want to set up a module for RNN, where function `__init__` sets up weight and biases, and function `forward` call the `rnn_forward` function from before.\n",
    "\n",
    "**We have written this part in `RNN` for you but you are highly recommended to go through the code as you will write `modules` on your own later.**\n",
    "\n",
    "All the implementation will be with `autograd` and `nn.Module` going forward. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CIjmnjRd6bhZ"
   },
   "source": [
    "# RNN for image captioning\n",
    "You will implement a few necessary tools and layers in order to build an image captioning model (class `CaptioningRNN`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IAa1Kvl2P_2k"
   },
   "source": [
    "## Image Feature Extraction\n",
    "Here, we use [MobileNet v2](https://pytorch.org/hub/pytorch_vision_mobilenet_v2/) for image feature extraction. For vanilla RNN and LSTM, we use the pooled CNN feature activation. For Attention LSTM, we use the CNN feature activation map after the last convolution layer. Checkout the `FeatureExtractor` method in `rnn_lstm_attention_captioning.py` to see the initialization of the model.\n",
    "\n",
    "Now, let's see what's inside MobileNet v2. Assume we have a 3x112x112 image input. We pass argument `pooling=True` to the model so the CNN activation is spatially-pooled from `1280x4x4` to `1280`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_pV0Lau_yDwX"
   },
   "outputs": [],
   "source": [
    "model = FeatureExtractor(pooling=True, verbose=True, device='cuda')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SVAxU-jO6bhR"
   },
   "source": [
    "## Word embedding\n",
    "In deep learning systems, we commonly represent words using vectors. Each word of the vocabulary will be associated with a vector, and these vectors will be learned jointly with the rest of the system.\n",
    "\n",
    "Implement the module `WordEmbedding.forward` to convert words (represented by integers) into vectors. Run the following to check your implementation. You should see an error on the order of `1e-7` or less. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BZuz2ieE6bhR"
   },
   "outputs": [],
   "source": [
    "N, T, V, D = 2, 4, 5, 3\n",
    "\n",
    "x = torch.tensor([[0, 3, 1, 2], [2, 1, 0, 3]], **to_long_cuda)\n",
    "W = torch.linspace(0, 1, steps=V*D, **to_double_cuda).reshape(V, D)\n",
    "\n",
    "# YOUR_TURN: Impelement WordEmbedding\n",
    "model_emb = WordEmbedding(V, D, **to_double_cuda)\n",
    "model_emb.W_embed.data.copy_(W)\n",
    "out = model_emb(x)\n",
    "expected_out = torch.tensor([\n",
    " [[ 0.,          0.07142857,  0.14285714],\n",
    "  [ 0.64285714,  0.71428571,  0.78571429],\n",
    "  [ 0.21428571,  0.28571429,  0.35714286],\n",
    "  [ 0.42857143,  0.5,         0.57142857]],\n",
    " [[ 0.42857143,  0.5,         0.57142857],\n",
    "  [ 0.21428571,  0.28571429,  0.35714286],\n",
    "  [ 0.,          0.07142857,  0.14285714],\n",
    "  [ 0.64285714,  0.71428571,  0.78571429]]], **to_double_cuda)\n",
    "\n",
    "print('out error: ', rel_error(expected_out, out))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6u-4Sr--6bhV",
    "tags": []
   },
   "source": [
    "## (Temporal) Affine layer\n",
    "At every timestep we use an affine function to transform the RNN hidden vector at that timestep into scores for each word in the vocabulary. This could be easily done with the [`nn.Linear`](https://pytorch.org/docs/master/nn.html#torch.nn.Linear) module. It can also work as a regular affine layer, like the one you have implemented from previous assignments. Run the following examples to see how it works. You will intensively use `nn.Linear` later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SA2DgU4X6bhV",
    "tags": []
   },
   "outputs": [],
   "source": [
    "reset_seed(0)\n",
    "\n",
    "N, T, D, M = 2, 3, 4, 3\n",
    "\n",
    "w = torch.linspace(-0.2, 0.4, steps=D*M, **to_double_cuda).reshape(D, M).permute(1, 0)\n",
    "b = torch.linspace(-0.4, 0.1, steps=M, **to_double_cuda)\n",
    "\n",
    "temporal_affine = nn.Linear(D, M).to(**to_double_cuda)\n",
    "temporal_affine.weight.data.copy_(w)\n",
    "temporal_affine.bias.data.copy_(b)\n",
    "\n",
    "# For regular affine layer\n",
    "x = torch.linspace(-0.1, 0.3, steps=N*D, **to_double_cuda).reshape(N, D)\n",
    "out = temporal_affine(x)\n",
    "print('affine layer - input shape: {}, output shape: {}'.format(x.shape, out.shape))\n",
    "correct_out = torch.tensor([[-0.35584416, -0.10896104,  0.13792208],\n",
    "                     [-0.31428571, -0.01753247,  0.27922078]], **to_double_cuda)\n",
    "\n",
    "print('dx error: ', rel_error(out, correct_out))\n",
    "\n",
    "\n",
    "# For temporal affine layer\n",
    "x = torch.linspace(-0.1, 0.3, steps=N*T*D, **to_double_cuda).reshape(N, T, D)\n",
    "out = temporal_affine(x)\n",
    "print('\\ntemporal affine layer - input shape: {}, output shape: {}'.format(x.shape, out.shape))\n",
    "correct_out = torch.tensor([[[-0.39920949, -0.16533597,  0.06853755],\n",
    "                             [-0.38656126, -0.13750988,  0.11154150],\n",
    "                             [-0.37391304, -0.10968379,  0.15454545]],\n",
    "                            [[-0.36126482, -0.08185771,  0.19754941],\n",
    "                             [-0.34861660, -0.05403162,  0.24055336],\n",
    "                             [-0.33596838, -0.02620553,  0.28355731]]], **to_double_cuda)\n",
    "\n",
    "print('dx error: ', rel_error(out, correct_out))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K6Py13Ak6bhX",
    "tags": []
   },
   "source": [
    "## Temporal Softmax loss\n",
    "In an RNN language model, at every timestep we produce a score for each word in the vocabulary. We know the ground-truth word at each timestep, so we use a softmax loss function to compute loss and gradient at each timestep. We sum the losses over time and average them over the minibatch.\n",
    "\n",
    "However there is one wrinkle: since we operate over minibatches and different captions may have different lengths, we append `<NULL>` tokens to the end of each caption so they all have the same length. We don't want these `<NULL>` tokens to count toward the loss or gradient, so in addition to scores and ground-truth labels our loss function also accepts a `ignore_index` that tells it which index in caption should be ignored when computing the loss.\n",
    "\n",
    "Implement the `temporal_softmax_loss` and run the follwing cell to check if the implmentation is correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nlFvgXtD6bhX",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def check_loss(N, T, V, p):\n",
    "    x = 0.001 * torch.randn(N, T, V, **to_double_cuda)\n",
    "    y = torch.randint(V, size=(N, T), **to_long_cuda)\n",
    "    mask = torch.rand(N, T, **to_double_cuda)\n",
    "    y[mask > p] = 0\n",
    "    # YOUR_TURN: Impelement temporal_softmax_loss\n",
    "    print(temporal_softmax_loss(x, y, NULL_index).item())\n",
    "  \n",
    "check_loss(1000, 1, 10, 1.0)   # Should be about 2.00-2.11\n",
    "check_loss(1000, 10, 10, 1.0)  # Should be about 20.6-21.0\n",
    "check_loss(5000, 10, 10, 0.1) # Should be about 2.00-2.11"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XWrmaSZaUxqX"
   },
   "source": [
    "## Captioning Module\n",
    "Now we are wrapping everything into the captioning module. Implement the `CaptioningRNN.__init__` function for initialization and the `CaptioningRNN.forward` for the forward pass. For now you only need to implement for the case where `cell_type='rnn'`, indicating vanialla RNNs; you will implement the `LSTM` case and `AttentionLSTM` case later.\n",
    "\n",
    "Skip the `CaptioningRNN.sample` for now, we will come back to this later.\n",
    "\n",
    "Run the following to check your forward pass using a small test case; you should see difference on the order of `1e-7` or less."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "d8a71FL_6bhZ"
   },
   "outputs": [],
   "source": [
    "reset_seed(0)\n",
    "\n",
    "N, D, W, H = 10, 1280, 30, 40\n",
    "D_img = 112\n",
    "word_to_idx = {'<NULL>': 0, 'cat': 2, 'dog': 3}\n",
    "V = len(word_to_idx)\n",
    "T = 13\n",
    "\n",
    "# YOUR_TURN: Impelement CaptioningRNN\n",
    "model = CaptioningRNN(word_to_idx,\n",
    "          input_dim=D,\n",
    "          wordvec_dim=W,\n",
    "          hidden_dim=H,\n",
    "          cell_type='rnn',\n",
    "          ignore_index=NULL_index,\n",
    "          **to_float_cuda) # use float here to be consistent with MobileNet v2\n",
    "\n",
    "\n",
    "for k,v in model.named_parameters():\n",
    "  # print(k, v.shape) # uncomment this to see the weight shape\n",
    "  v.data.copy_(torch.linspace(-1.4, 1.3, steps=v.numel()).reshape(*v.shape))\n",
    "\n",
    "images = torch.linspace(-3., 3., steps=(N * 3 * D_img * D_img),\n",
    "                       **to_float_cuda).reshape(N, 3, D_img, D_img)\n",
    "captions = (torch.arange(N * T, **to_long_cuda) % V).reshape(N, T)\n",
    "\n",
    "loss = model(images, captions).item()\n",
    "expected_loss = 150.6090393066\n",
    "\n",
    "print('loss: ', loss)\n",
    "print('expected loss: ', expected_loss)\n",
    "print('difference: ', rel_error(torch.tensor(loss), torch.tensor(expected_loss)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7YAOcQ4h6bhc"
   },
   "source": [
    "## Image Captioning solver\n",
    "Different from the `Solver` class that we used to train image classification models on the previous assignment, on this assignment we use the `torch.optim` package to train image captioning models.\n",
    "\n",
    "We have written this part for you and you need to train the model and generate plots on the training loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bRC0lNgOxheT"
   },
   "outputs": [],
   "source": [
    "def captioning_train(rnn_model, image_data, caption_data, lr_decay=1, **kwargs):\n",
    "  \"\"\"\n",
    "  Run optimization to train the model.\n",
    "  \"\"\"\n",
    "  # optimizer setup\n",
    "  from torch import optim\n",
    "  optimizer = optim.Adam(\n",
    "    filter(lambda p: p.requires_grad, rnn_model.parameters()),\n",
    "    learning_rate) # leave betas and eps by default\n",
    "  lr_scheduler = optim.lr_scheduler.LambdaLR(optimizer,\n",
    "                                             lambda epoch: lr_decay ** epoch)\n",
    "\n",
    "  # sample minibatch data\n",
    "  iter_per_epoch = math.ceil(image_data.shape[0] // batch_size)\n",
    "  loss_history = []\n",
    "  rnn_model.train()\n",
    "  for i in range(num_epochs):\n",
    "    start_t = time.time()\n",
    "    for j in range(iter_per_epoch):\n",
    "      images, captions = image_data[j*batch_size:(j+1)*batch_size], \\\n",
    "                           caption_data[j*batch_size:(j+1)*batch_size]\n",
    "\n",
    "      loss = rnn_model(images, captions)\n",
    "      optimizer.zero_grad()\n",
    "      loss.backward()\n",
    "      loss_history.append(loss.item())\n",
    "      optimizer.step()\n",
    "    end_t = time.time()\n",
    "    print('(Epoch {} / {}) loss: {:.4f} time per epoch: {:.1f}s'.format(\n",
    "        i, num_epochs, loss.item(), end_t-start_t))\n",
    "\n",
    "    lr_scheduler.step()\n",
    "\n",
    "  # plot the training losses\n",
    "  plt.plot(loss_history)\n",
    "  plt.xlabel('Iteration')\n",
    "  plt.ylabel('Loss')\n",
    "  plt.title('Training loss history')\n",
    "  plt.show()\n",
    "  return rnn_model, loss_history"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eUc5c217pFYE"
   },
   "source": [
    "## Overfit small data\n",
    "Once you have familiarized yourself with the `optim` API used above, run the following to make sure your model overfits a small sample of 50 training examples. You should see a final loss of less than `1.6` and it should be done fairly quickly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yzhsGRzk6bhd"
   },
   "outputs": [],
   "source": [
    "reset_seed(0)\n",
    "\n",
    "# data input\n",
    "small_num_train = 50\n",
    "sample_idx = torch.linspace(0, num_train-1, steps=small_num_train, **to_long_cuda)\n",
    "small_image_data = data_dict['train_images'][sample_idx].to('cuda')\n",
    "small_caption_data = data_dict['train_captions'][sample_idx].to('cuda')\n",
    "\n",
    "# optimization arguments\n",
    "num_epochs = 80\n",
    "batch_size = 50\n",
    "  \n",
    "# create the image captioning model\n",
    "model = CaptioningRNN(\n",
    "          cell_type='rnn',\n",
    "          word_to_idx=data_dict['vocab']['token_to_idx'],\n",
    "          input_dim=1280, # hard-coded, do not modify\n",
    "          hidden_dim=512,\n",
    "          wordvec_dim=256,\n",
    "          ignore_index=NULL_index,\n",
    "          **to_float_cuda)\n",
    "\n",
    "for learning_rate in [1e-3]:\n",
    "  print('learning rate is: ', learning_rate)\n",
    "  rnn_overfit, _ = captioning_train(model, small_image_data, small_caption_data,\n",
    "                num_epochs=num_epochs, batch_size=batch_size,\n",
    "                learning_rate=learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UiHsRysE6bhe"
   },
   "source": [
    "## Caption sampling\n",
    "Unlike classification models, image captioning models behave very differently at training time and at test time. At training time, we have access to the ground-truth caption, so we feed ground-truth words as input to the RNN at each timestep. At test time, we sample from the distribution over the vocabulary at each timestep, and feed the sample as input to the RNN at the next timestep.\n",
    "\n",
    "Implement the `CaptioningRNN.sample` for test-time sampling. After doing so, run the following to train a captioning model and sample from the model on both training and validation data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BN_sn1YcFZz4"
   },
   "source": [
    "### Train the net\n",
    "After you are done implementing the `CaptioningRNN.sample` method, perform the training on the entire training set. You should see a final loss less than `2.0` and each epoch should take ~6s to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dXHnPuM_FU7k"
   },
   "outputs": [],
   "source": [
    "reset_seed(0)\n",
    "\n",
    "# data input\n",
    "small_num_train = num_train\n",
    "sample_idx = torch.randint(num_train, size=(small_num_train,), **to_long_cuda)\n",
    "small_image_data = data_dict['train_images'][sample_idx].to('cuda')\n",
    "small_caption_data = data_dict['train_captions'][sample_idx].to('cuda')\n",
    "\n",
    "# optimization arguments\n",
    "num_epochs = 60\n",
    "batch_size = 250\n",
    "\n",
    "# create the image captioning model\n",
    "rnn_model = CaptioningRNN(\n",
    "          cell_type='rnn',\n",
    "          word_to_idx=data_dict['vocab']['token_to_idx'],\n",
    "          input_dim=1280, # hard-coded, do not modify\n",
    "          hidden_dim=512,\n",
    "          wordvec_dim=256,\n",
    "          ignore_index=NULL_index,\n",
    "          **to_float_cuda)\n",
    "\n",
    "for learning_rate in [1e-3]:\n",
    "  print('learning rate is: ', learning_rate)\n",
    "  rnn_model_submit, rnn_loss_submit = captioning_train(rnn_model, small_image_data, small_caption_data,\n",
    "                num_epochs=num_epochs, batch_size=batch_size,\n",
    "                learning_rate=learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "97xga3Q5GO8B"
   },
   "source": [
    "### Test-time sampling\n",
    "The samples on training data should be very good; the samples on validation data will probably make less sense."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Rvt326nX6bhf"
   },
   "outputs": [],
   "source": [
    "# Sample a minibatch and show the reshaped 112x112 images,\n",
    "# GT captions, and generated captions by your model.\n",
    "batch_size = 3\n",
    "\n",
    "for split in ['train', 'val']:\n",
    "  sample_idx = torch.randint(0, num_train if split=='train' else num_val, (batch_size,))\n",
    "  sample_images = data_dict[split+'_images'][sample_idx]\n",
    "  sample_captions = data_dict[split+'_captions'][sample_idx]\n",
    "\n",
    "  # decode_captions is loaded from a4_helper.py\n",
    "  gt_captions = decode_captions(sample_captions, data_dict['vocab']['idx_to_token'])\n",
    "  rnn_model.eval()\n",
    "  generated_captions = rnn_model.sample(sample_images)\n",
    "  generated_captions = decode_captions(generated_captions, data_dict['vocab']['idx_to_token'])\n",
    "\n",
    "  for i in range(batch_size):\n",
    "    plt.imshow(sample_images[i].permute(1, 2, 0))\n",
    "    plt.axis('off')\n",
    "    plt.title('%s\\nRNN Generated:%s\\nGT:%s' % (split, generated_captions[i], gt_captions[i]))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vVeGWWQTvMMF",
    "tags": [
     "pdf-title"
    ]
   },
   "source": [
    "# Image Captioning with LSTMs\n",
    "In the previous exercise you implemented a vanilla RNN and applied it to image captioning. Next, we will implement the LSTM update rule and use it for image captioning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_8Zd6FGPvMMa"
   },
   "source": [
    "# LSTM\n",
    "If you read recent papers, you'll see that many people use a variant on the vanilla RNN called Long-Short Term Memory (LSTM) RNNs. Vanilla RNNs can be tough to train on long sequences due to vanishing and exploding gradients caused by repeated matrix multiplication. LSTMs solve this problem by replacing the simple update rule of the vanilla RNN with a gating mechanism as follows.\n",
    "\n",
    "Similar to the vanilla RNN, at each timestep we receive an input $x_t\\in\\mathbb{R}^D$ and the previous hidden state $h_{t-1}\\in\\mathbb{R}^H$; the LSTM also maintains an $H$-dimensional *cell state*, so we also receive the previous cell state $c_{t-1}\\in\\mathbb{R}^H$. The learnable parameters of the LSTM are an *input-to-hidden* matrix $W_x\\in\\mathbb{R}^{4H\\times D}$, a *hidden-to-hidden* matrix $W_h\\in\\mathbb{R}^{4H\\times H}$ and a *bias vector* $b\\in\\mathbb{R}^{4H}$.\n",
    "\n",
    "At each timestep we first compute an *activation vector* $a\\in\\mathbb{R}^{4H}$ as $a=W_xx_t + W_hh_{t-1}+b$. We then divide this into four vectors $a_i,a_f,a_o,a_g\\in\\mathbb{R}^H$ where $a_i$ consists of the first $H$ elements of $a$, $a_f$ is the next $H$ elements of $a$, etc. We then compute the *input gate* $g\\in\\mathbb{R}^H$, *forget gate* $f\\in\\mathbb{R}^H$, *output gate* $o\\in\\mathbb{R}^H$ and *block input* $g\\in\\mathbb{R}^H$ as\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "i = \\sigma(a_i) \\hspace{2pc}\n",
    "f = \\sigma(a_f) \\hspace{2pc}\n",
    "o = \\sigma(a_o) \\hspace{2pc}\n",
    "g = \\tanh(a_g)\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "where $\\sigma$ is the sigmoid function and $\\tanh$ is the hyperbolic tangent, both applied elementwise.\n",
    "\n",
    "Finally we compute the next cell state $c_t$ and next hidden state $h_t$ as\n",
    "\n",
    "$$\n",
    "c_{t} = f\\odot c_{t-1} + i\\odot g \\hspace{4pc}\n",
    "h_t = o\\odot\\tanh(c_t)\n",
    "$$\n",
    "\n",
    "where $\\odot$ is the elementwise product of vectors.\n",
    "\n",
    "In the rest of the notebook we will implement the LSTM update rule and apply it to the image captioning task. \n",
    "\n",
    "In the code, we assume that data is stored in batches so that $X_t \\in \\mathbb{R}^{N\\times D}$, and will work with *transposed* versions of the parameters: $W_x \\in \\mathbb{R}^{D \\times 4H}$, $W_h \\in \\mathbb{R}^{H\\times 4H}$ so that activations $A \\in \\mathbb{R}^{N\\times 4H}$ can be computed efficiently as $A = X_t W_x + H_{t-1} W_h$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t4DNkZYevMMc"
   },
   "source": [
    "## LSTM: step forward\n",
    "Implement the forward pass for a single timestep of an LSTM in the `lstm_step_forward` function. This should be similar to the `rnn_step_forward` function that you implemented above, but using the LSTM update rule instead.\n",
    "\n",
    "Don't worry about the backward part! `autograd` will handle it.\n",
    "\n",
    "Once you are done, run the following to perform a simple test of your implementation. You should see errors on the order of `1e-7` or less."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "awLF_A5ZvMMd"
   },
   "outputs": [],
   "source": [
    "N, D, H = 3, 4, 5\n",
    "x = torch.linspace(-0.4, 1.2, steps=N*D, **to_double_cuda).reshape(N, D)\n",
    "prev_h = torch.linspace(-0.3, 0.7, steps=N*H, **to_double_cuda).reshape(N, H)\n",
    "prev_c = torch.linspace(-0.4, 0.9, steps=N*H, **to_double_cuda).reshape(N, H)\n",
    "Wx = torch.linspace(-2.1, 1.3, steps=4*D*H, **to_double_cuda).reshape(D, 4 * H)\n",
    "Wh = torch.linspace(-0.7, 2.2, steps=4*H*H, **to_double_cuda).reshape(H, 4 * H)\n",
    "b = torch.linspace(0.3, 0.7, steps=4*H, **to_double_cuda)\n",
    "\n",
    "# YOUR_TURN: Impelement lstm_step_forward\n",
    "next_h, next_c = lstm_step_forward(x, prev_h, prev_c, Wx, Wh, b)\n",
    "\n",
    "expected_next_h = torch.tensor([\n",
    "    [ 0.24635157,  0.28610883,  0.32240467,  0.35525807,  0.38474904],\n",
    "    [ 0.49223563,  0.55611431,  0.61507696,  0.66844003,  0.7159181 ],\n",
    "    [ 0.56735664,  0.66310127,  0.74419266,  0.80889665,  0.858299  ]], **to_double_cuda)\n",
    "expected_next_c = torch.tensor([\n",
    "    [ 0.32986176,  0.39145139,  0.451556,    0.51014116,  0.56717407],\n",
    "    [ 0.66382255,  0.76674007,  0.87195994,  0.97902709,  1.08751345],\n",
    "    [ 0.74192008,  0.90592151,  1.07717006,  1.25120233,  1.42395676]], **to_double_cuda)\n",
    "\n",
    "print('next_h error: ', rel_error(expected_next_h, next_h))\n",
    "print('next_c error: ', rel_error(expected_next_c, next_c))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ErgRQwwzvMMt"
   },
   "source": [
    "## LSTM: forward\n",
    "Implement the `lstm_forward` function to run an LSTM forward on an entire timeseries of data.\n",
    "\n",
    "Again, don't worry about the backward part! `autograd` will handle it.\n",
    "\n",
    "When you are done, run the following to check your implementation. You should see an error on the order of `1e-7` or less."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_x-3BJiEvMMv"
   },
   "outputs": [],
   "source": [
    "N, D, H, T = 2, 5, 4, 3\n",
    "x = torch.linspace(-0.4, 0.6, steps=N*T*D, **to_double_cuda).reshape(N, T, D)\n",
    "h0 = torch.linspace(-0.4, 0.8, steps=N*H, **to_double_cuda).reshape(N, H)\n",
    "Wx = torch.linspace(-0.2, 0.9, steps=4*D*H, **to_double_cuda).reshape(D, 4 * H)\n",
    "Wh = torch.linspace(-0.3, 0.6, steps=4*H*H, **to_double_cuda).reshape(H, 4 * H)\n",
    "b = torch.linspace(0.2, 0.7, steps=4*H, **to_double_cuda)\n",
    "\n",
    "# YOUR_TURN: Impelement lstm_forward\n",
    "h = lstm_forward(x, h0, Wx, Wh, b)\n",
    "\n",
    "expected_h = torch.tensor([\n",
    " [[ 0.01764008,  0.01823233,  0.01882671,  0.0194232 ],\n",
    "  [ 0.11287491,  0.12146228,  0.13018446,  0.13902939],\n",
    "  [ 0.31358768,  0.33338627,  0.35304453,  0.37250975]],\n",
    " [[ 0.45767879,  0.4761092,   0.4936887,   0.51041945],\n",
    "  [ 0.6704845,   0.69350089,  0.71486014,  0.7346449 ],\n",
    "  [ 0.81733511,  0.83677871,  0.85403753,  0.86935314]]], **to_double_cuda)\n",
    "\n",
    "print('h error: ', rel_error(expected_h, h))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_YZ3OqrqkjLt"
   },
   "source": [
    "## LSTM Module\n",
    "\n",
    "We can now wrap the LSTM functions we wrote into an nn.Module. You can check the `LSTM` module in the script."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "92SStL_tvMNK"
   },
   "source": [
    "## LSTM captioning model\n",
    "\n",
    "Now that you have implemented an LSTM, update the `CaptioningRNN.__init__` and `CaptioningRNN.forward` implementation method **ONLY** to also handle the case where `self.cell_type` is `lstm`. **This should require adding less than 5 lines of code.**\n",
    "\n",
    "Once you have done so, run the following to check your implementation. You should see a difference on the order of `1e-7` or less."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NNpiC4WSvMNL"
   },
   "outputs": [],
   "source": [
    "reset_seed(0)\n",
    "\n",
    "N, D, W, H = 10, 1280, 30, 40\n",
    "D_img = 112\n",
    "word_to_idx = {'<NULL>': 0, 'cat': 2, 'dog': 3}\n",
    "V = len(word_to_idx)\n",
    "T = 13\n",
    "\n",
    "# YOUR_TURN: Impelement CaptioningRNN for lstm \n",
    "model = CaptioningRNN(word_to_idx,\n",
    "          input_dim=D,\n",
    "          wordvec_dim=W,\n",
    "          hidden_dim=H,\n",
    "          cell_type='lstm',\n",
    "          ignore_index=NULL_index,\n",
    "          **to_float_cuda)\n",
    "\n",
    "for k,v in model.named_parameters():\n",
    "  # print(k, v.shape) # uncomment this to see the weight shape\n",
    "  v.data.copy_(torch.linspace(-1.4, 1.3, steps=v.numel()).reshape(*v.shape))\n",
    "\n",
    "images = torch.linspace(-3., 3., steps=(N * 3 * D_img * D_img),\n",
    "                       **to_float_cuda).reshape(N, 3, D_img, D_img)\n",
    "captions = (torch.arange(N * T, **to_long_cuda) % V).reshape(N, T)\n",
    "\n",
    "loss = model(images, captions).item()\n",
    "expected_loss = 146.3161468505\n",
    "\n",
    "print('loss: ', loss)\n",
    "print('expected loss: ', expected_loss)\n",
    "print('difference: ', rel_error(torch.tensor(loss), torch.tensor(expected_loss)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "06hbDnRXvMNO"
   },
   "source": [
    "## Overfit small data\n",
    "We have written this part for you. Run the following to overfit an LSTM captioning model on the same small dataset as we used for the RNN previously. You should see a final loss less than `4` after 80 epochs and it should run fairly quickly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "O-tETnd3vMNP"
   },
   "outputs": [],
   "source": [
    "reset_seed(0)\n",
    "\n",
    "# data input\n",
    "small_num_train = 50\n",
    "sample_idx = torch.linspace(0, num_train-1, steps=small_num_train, **to_float_cuda).long()\n",
    "small_image_data = data_dict['train_images'][sample_idx].to('cuda')\n",
    "small_caption_data = data_dict['train_captions'][sample_idx].to('cuda')\n",
    "\n",
    "# optimization arguments\n",
    "num_epochs = 80\n",
    "batch_size = 50\n",
    "\n",
    "# create the image captioning model\n",
    "model = CaptioningRNN(\n",
    "          cell_type='lstm',\n",
    "          word_to_idx=data_dict['vocab']['token_to_idx'],\n",
    "          input_dim=1280, # hard-coded, do not modify\n",
    "          hidden_dim=512,\n",
    "          wordvec_dim=256,\n",
    "          ignore_index=NULL_index,\n",
    "          **to_float_cuda)\n",
    "\n",
    "for learning_rate in [1e-2]:\n",
    "  print('learning rate is: ', learning_rate)\n",
    "  lstm_overfit, _ = captioning_train(model, small_image_data, small_caption_data,\n",
    "                num_epochs=num_epochs, batch_size=batch_size,\n",
    "                learning_rate=learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4vzLUzlWvMNT"
   },
   "source": [
    "## Caption sampling\n",
    "Modify the  `CaptioningRNN.sample` method in class to handle the case where `self.cell_type` is `lstm`. **This should take fewer than 10 lines of code.**\n",
    "\n",
    "When you are done, run the following to train a captioning model and sample from your the model on some training and validation set samples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bn2PpJL5oC0J"
   },
   "source": [
    "### Train the net\n",
    "Now, perform the training on the entire training set. You should see a final loss less than `2.8`. Each epoch should take ~7s to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "f9MFRowdoHW7"
   },
   "outputs": [],
   "source": [
    "reset_seed(0)\n",
    "\n",
    "# data input\n",
    "small_num_train = num_train\n",
    "sample_idx = torch.randint(num_train, size=(small_num_train,), **to_long_cuda)\n",
    "small_image_data = data_dict['train_images'][sample_idx].to('cuda')\n",
    "small_caption_data = data_dict['train_captions'][sample_idx].to('cuda')\n",
    "\n",
    "# optimization arguments\n",
    "num_epochs = 60\n",
    "batch_size = 250\n",
    "\n",
    "# create the image captioning model\n",
    "lstm_model = CaptioningRNN(\n",
    "          cell_type='lstm',\n",
    "          word_to_idx=data_dict['vocab']['token_to_idx'],\n",
    "          input_dim=1280, # hard-coded, do not modify\n",
    "          hidden_dim=512,\n",
    "          wordvec_dim=256,\n",
    "          ignore_index=NULL_index,\n",
    "          **to_float_cuda)\n",
    "\n",
    "for learning_rate in [1e-3]:\n",
    "  print('learning rate is: ', learning_rate)\n",
    "  lstm_model_submit, lstm_loss_submit = captioning_train(lstm_model, small_image_data, small_caption_data,\n",
    "                num_epochs=num_epochs, batch_size=batch_size,\n",
    "                learning_rate=learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wsM2pIYpG3v1"
   },
   "source": [
    "### Test-time sampling\n",
    "As with the RNN, the samples on training data should be very good; the samples on validation data will probably make less sense."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ziQJ7SBnvMNU"
   },
   "outputs": [],
   "source": [
    "# Sample a minibatch and show the reshaped 112x112 images,\n",
    "# GT captions, and generated captions by your model.\n",
    "batch_size = 3\n",
    "\n",
    "for split in ['train', 'val']:\n",
    "  sample_idx = torch.randint(0, num_train if split=='train' else num_val, (batch_size,))\n",
    "  sample_images = data_dict[split+'_images'][sample_idx]\n",
    "  sample_captions = data_dict[split+'_captions'][sample_idx]\n",
    "\n",
    "  # decode_captions is loaded from a4_helper.py\n",
    "  gt_captions = decode_captions(sample_captions, data_dict['vocab']['idx_to_token'])\n",
    "  lstm_model.eval()\n",
    "  generated_captions = lstm_model.sample(sample_images)\n",
    "  generated_captions = decode_captions(generated_captions, data_dict['vocab']['idx_to_token'])\n",
    "\n",
    "  for i in range(batch_size):\n",
    "    plt.imshow(sample_images[i].permute(1, 2, 0))\n",
    "    plt.axis('off')\n",
    "    plt.title('%s\\nLSTM Generated:%s\\nGT:%s' % (split, generated_captions[i], gt_captions[i]))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ECcPPE_Pqc8v"
   },
   "source": [
    "# Attention LSTM\n",
    "Attention LSTM essentially adds an attention input $x_{attn}^t\\in\\mathbb{R}^H$ into LSTM, along with $x_t\\in\\mathbb{R}^D$ and the previous hidden state $h_{t-1}\\in\\mathbb{R}^H$.\n",
    "\n",
    "To get the attention input $x_{attn}^t$, here we adopt a method called `scaled dot-product attention`, as covered in the lecture. We first project the CNN feature activation from $\\mathbb{R}^{1280\\times4\\times4}$ to $\\mathbb{R}^{H\\times4\\times4}$ using an affine layer. Given the projected activation $A\\in \\mathbb{R}^{H\\times4\\times4}$ and the LSTM hidden state from the previous time step $h_{t-1}$, we formuate the attention weights on $A$ at time step $t$ as $M_{attn}^t=h_{t-1}A/\\sqrt{H} \\in \\mathbb{R}^{4\\times4}$.\n",
    "\n",
    "To simplify the formuation here, we flatten the spatial dimensions of $A$ and $M_{attn}^t$ which gives $\\tilde{A}\\in \\mathbb{R}^{H\\times16}$ and $\\tilde{M^t}_{attn}=h_{t-1}A\\in \\mathbb{R}^{16}$.\n",
    "We add a **`softmax`** activation function on $\\tilde{M^t}_{attn}$ so that the attention weights at each time step are normalized and sum up to one.\n",
    "\n",
    "The attention embedding given the attention weights is then $x_{attn}^t=\\tilde{A}\\tilde{M^t}_{attn} \\in\\mathbb{R}^H$.\n",
    "\n",
    "You will implement a batch version of the attention layer we have described here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GTDk54Q4ubQ1"
   },
   "source": [
    "## Scaled dot-product attention\n",
    "Implement the `dot_product_attention` function. Given the LSTM hidden state from the previous time step `prev_h` (or $h_{t-1}$) and the projected CNN feature activation `A`, compute the attention weights `attn_weights` (or $\\tilde{M^t}_{attn}$ with a reshaping to $\\mathbb{R}^{4\\times4}$) attention embedding output `attn` (or $x_{attn}^t$) using the formulation we provided.\n",
    "\n",
    "When you are done, run the following to check your implementation. You should see an error on the order of `1e-7` or less."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "irAslXWfaVGw"
   },
   "outputs": [],
   "source": [
    "N, H = 2, 5\n",
    "D_a = 4\n",
    "\n",
    "prev_h = torch.linspace(-0.4, 0.6, steps=N*H, **to_double_cuda).reshape(N, H)\n",
    "A = torch.linspace(-0.4, 1.8, steps=N*H*D_a*D_a, **to_double_cuda).reshape(N, H, D_a, D_a)\n",
    "\n",
    "# YOUR_TURN: Impelement dot_product_attention\n",
    "attn, attn_weights = dot_product_attention(prev_h, A)\n",
    "\n",
    "expected_attn = torch.tensor([[-0.29784344, -0.07645979,  0.14492386,  0.36630751,  0.58769115],\n",
    "        [ 0.81412643,  1.03551008,  1.25689373,  1.47827738,  1.69966103]], **to_double_cuda)\n",
    "expected_attn_weights = torch.tensor([[[0.06511126, 0.06475411, 0.06439892, 0.06404568],\n",
    "         [0.06369438, 0.06334500, 0.06299754, 0.06265198],\n",
    "         [0.06230832, 0.06196655, 0.06162665, 0.06128861],\n",
    "         [0.06095243, 0.06061809, 0.06028559, 0.05995491]],\n",
    "\n",
    "        [[0.05717142, 0.05784357, 0.05852362, 0.05921167],\n",
    "         [0.05990781, 0.06061213, 0.06132473, 0.06204571],\n",
    "         [0.06277517, 0.06351320, 0.06425991, 0.06501540],\n",
    "         [0.06577977, 0.06655312, 0.06733557, 0.06812722]]], **to_double_cuda)\n",
    "\n",
    "print('attn error: ', rel_error(expected_attn, attn))\n",
    "print('attn_weights error: ', rel_error(expected_attn_weights, attn_weights))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2VK_Ixn1qlRo"
   },
   "source": [
    "## Attention LSTM: step forward"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DVddQlj4xwRk"
   },
   "source": [
    "Modify the `lstm_step_forward` function from earlier to support the extra attention input `attn` (or $x_{attn}$) and its embedding weight matrix `Wattn` (or $W_{attn}$) in the LSTM cell. Hence, at each timestep the *activation vector* $a\\in\\mathbb{R}^{4H}$ in LSTM cell is formulated as:\n",
    "\n",
    "$a=W_xx_t + W_hh_{t-1}+W_{attn}x_{attn}^t+b$.\n",
    "\n",
    "\n",
    "**This should require adding less than 5 lines of code.**\n",
    "\n",
    "Once you are done, run the following to perform a simple test of your implementation. You should see errors on the order of `1e-8` or less."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oaS31Ncf3l0d"
   },
   "outputs": [],
   "source": [
    "N, D, H = 3, 4, 5\n",
    "\n",
    "x = torch.linspace(-0.4, 1.2, steps=N*D, **to_double_cuda).reshape(N, D)\n",
    "prev_h = torch.linspace(-0.3, 0.7, steps=N*H, **to_double_cuda).reshape(N, H)\n",
    "prev_c = torch.linspace(-0.4, 0.9, steps=N*H, **to_double_cuda).reshape(N, H)\n",
    "Wx = torch.linspace(-2.1, 1.3, steps=4*D*H, **to_double_cuda).reshape(D, 4 * H)\n",
    "Wh = torch.linspace(-0.7, 2.2, steps=4*H*H, **to_double_cuda).reshape(H, 4 * H)\n",
    "b = torch.linspace(0.3, 0.7, steps=4*H, **to_double_cuda)\n",
    "attn = torch.linspace(0.6, 1.8, steps=N*H, **to_double_cuda).reshape(N, H)\n",
    "Wattn = torch.linspace(1.3, 4.2, steps=4*H*H, **to_double_cuda).reshape(H, 4 * H)\n",
    "\n",
    "# YOUR_TURN: Impelement lstm_step_forward\n",
    "next_h, next_c = lstm_step_forward(x, prev_h, prev_c, Wx, Wh, b, attn, Wattn)\n",
    "\n",
    "expected_next_h = torch.tensor([\n",
    "    [0.53704256, 0.59980774, 0.65596820, 0.70569729, 0.74932626],\n",
    "    [0.78729857, 0.82010653, 0.84828362, 0.87235677, 0.89283167],\n",
    "    [0.91017981, 0.92483119, 0.93717126, 0.94754073, 0.95623746]], **to_double_cuda)\n",
    "expected_next_c = torch.tensor([\n",
    "    [0.59999328, 0.69285041, 0.78570758, 0.87856479, 0.97142202],\n",
    "    [1.06428558, 1.15714276, 1.24999992, 1.34285708, 1.43571424],\n",
    "    [1.52857143, 1.62142857, 1.71428571, 1.80714286, 1.90000000]], **to_double_cuda)\n",
    "\n",
    "print('next_h error: ', rel_error(expected_next_h, next_h))\n",
    "print('next_c error: ', rel_error(expected_next_c, next_c))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VS0JHfJ53agv"
   },
   "source": [
    "## Attention LSTM: forward\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QRry6hvQ7ywx"
   },
   "source": [
    "Now, implement the `attention_forward` function to run an Attention LSTM forward on an entire timeseries of data. You will have to use the `dot_product_attention` function and the `lstm_step_forward` function you implemented.\n",
    "\n",
    "Again, don't worry about the backward part! `autograd` will handle it.\n",
    "\n",
    "When you are done, run the following to check your implementation. You should see an error on the order of `1e-8` or less.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aB6VU8nl4SmS"
   },
   "outputs": [],
   "source": [
    "N, D, H, T = 2, 5, 4, 3\n",
    "D_a = 4\n",
    "\n",
    "x = torch.linspace(-0.4, 0.6, steps=N*T*D, **to_double_cuda).reshape(N, T, D)\n",
    "A = torch.linspace(-0.4, 1.8, steps=N*H*D_a*D_a, **to_double_cuda).reshape(N, H, D_a, D_a)\n",
    "Wx = torch.linspace(-0.2, 0.9, steps=4*D*H, **to_double_cuda).reshape(D, 4 * H)\n",
    "Wh = torch.linspace(-0.3, 0.6, steps=4*H*H, **to_double_cuda).reshape(H, 4 * H)\n",
    "Wattn = torch.linspace(1.3, 4.2, steps=4*H*H, **to_double_cuda).reshape(H, 4 * H)\n",
    "b = torch.linspace(0.2, 0.7, steps=4*H, **to_double_cuda)\n",
    "\n",
    "# YOUR_TURN: Impelement attention_forward\n",
    "h = attention_forward(x, A, Wx, Wh, Wattn, b)\n",
    "\n",
    "expected_h = torch.tensor([\n",
    "        [[0.56141729, 0.70274849, 0.80000386, 0.86349400],\n",
    "         [0.89556391, 0.92856726, 0.94950579, 0.96281018],\n",
    "         [0.96792077, 0.97535465, 0.98039623, 0.98392994]],\n",
    "\n",
    "        [[0.95065880, 0.97135490, 0.98344373, 0.99045552],\n",
    "         [0.99317679, 0.99607466, 0.99774317, 0.99870293],\n",
    "         [0.99907382, 0.99946784, 0.99969426, 0.99982435]]], **to_double_cuda)\n",
    "\n",
    "print('h error: ', rel_error(expected_h, h))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jGdYDBjDqofZ"
   },
   "source": [
    "## Attention LSTM Module\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9VzpyHuX6Jzc"
   },
   "source": [
    "## Attention LSTM captioning model\n",
    "\n",
    "We can now wrap the Attention LSTM functions we wrote into an nn.Module. You can look at the `AttentionLSTM` class for detailed implementation.\n",
    "\n",
    "Now that you have implemented an attention module, update the implementation of `CaptioningRNN.__init__` and `CaptioningRNN.forward` method to also handle the case where `self.cell_type` is `attention`. **This should require adding less than 10 lines of code.**\n",
    "\n",
    "Once you have done so, run the following to check your implementation. You should see a difference on the order of `1e-7` or less."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7VqGqDYw6Jzd"
   },
   "outputs": [],
   "source": [
    "reset_seed(0)\n",
    "\n",
    "N, D, W, H = 10, 1280, 30, 40\n",
    "D_img = 112\n",
    "word_to_idx = {'<NULL>': 0, 'cat': 2, 'dog': 3}\n",
    "V = len(word_to_idx)\n",
    "T = 13\n",
    "\n",
    "# YOUR_TURN: Modify CaptioningRNN for attention\n",
    "model = CaptioningRNN(word_to_idx,\n",
    "          input_dim=D,\n",
    "          wordvec_dim=W,\n",
    "          hidden_dim=H,\n",
    "          cell_type='attention',\n",
    "          ignore_index=NULL_index,\n",
    "          **to_float_cuda)\n",
    "\n",
    "for k,v in model.named_parameters():\n",
    "  # print(k, v.shape) # uncomment this to see the weight shape\n",
    "  v.data.copy_(torch.linspace(-1.4, 1.3, steps=v.numel()).reshape(*v.shape))\n",
    "\n",
    "images = torch.linspace(-3., 3., steps=(N * 3 * D_img * D_img),\n",
    "                       **to_float_cuda).reshape(N, 3, D_img, D_img)\n",
    "captions = (torch.arange(N * T, **to_long_cuda) % V).reshape(N, T)\n",
    "\n",
    "loss = model(images, captions).item()\n",
    "expected_loss = 46.9113769531\n",
    "\n",
    "print('loss: ', loss)\n",
    "print('expected loss: ', expected_loss)\n",
    "print('difference: ', rel_error(torch.tensor(loss), torch.tensor(expected_loss)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eYxXTAn4q0wV"
   },
   "source": [
    "## Overfit small data\n",
    "We have written this part for you. Run the following to overfit an Attention LSTM captioning model on the same small dataset as we used for the RNN previously. You should see a final loss less than `9`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tlK7lKUgWeDS"
   },
   "outputs": [],
   "source": [
    "reset_seed(0)\n",
    "\n",
    "# data input\n",
    "small_num_train = 50\n",
    "sample_idx = torch.linspace(0, num_train-1, steps=small_num_train, **to_float_cuda).long()\n",
    "small_image_data = data_dict['train_images'][sample_idx].to('cuda')\n",
    "small_caption_data = data_dict['train_captions'][sample_idx].to('cuda')\n",
    "\n",
    "# optimization arguments\n",
    "num_epochs = 80\n",
    "batch_size = 50\n",
    "\n",
    "# create the image captioning model\n",
    "model = CaptioningRNN(\n",
    "          cell_type='attention',\n",
    "          word_to_idx=data_dict['vocab']['token_to_idx'],\n",
    "          input_dim=1280, # hard-coded, do not modify\n",
    "          hidden_dim=512,\n",
    "          wordvec_dim=256,\n",
    "          ignore_index=NULL_index,\n",
    "          **to_float_cuda)\n",
    "\n",
    "for learning_rate in [1e-3]:\n",
    "  print('learning rate is: ', learning_rate)\n",
    "  attn_overfit, _ = captioning_train(model, small_image_data, small_caption_data,\n",
    "                num_epochs=num_epochs, batch_size=batch_size,\n",
    "                learning_rate=learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ircMb7_qq7vB"
   },
   "source": [
    "## Caption sampling\n",
    "Modify the `CaptioningRNN.sample` method to handle the case where `self.cell_type` is `attention`. **This should take fewer than 10 lines of code.**\n",
    "\n",
    "When you are done run the following to train a captioning model and sample from the model on some training and validation set samples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t--pa33Sq4SW"
   },
   "source": [
    "### Train the net\n",
    "Now, perform the training on the entire training set. You should see a final loss less than `1.0`. Each epoch should take ~8s to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ScBvAfcXdVv4"
   },
   "outputs": [],
   "source": [
    "reset_seed(0)\n",
    "\n",
    "# data input\n",
    "small_num_train = num_train\n",
    "sample_idx = torch.randint(num_train, size=(small_num_train,), **to_long_cuda)\n",
    "small_image_data = data_dict['train_images'][sample_idx].to('cuda')\n",
    "small_caption_data = data_dict['train_captions'][sample_idx].to('cuda')\n",
    "\n",
    "# optimization arguments\n",
    "num_epochs = 60\n",
    "batch_size = 250\n",
    "\n",
    "# create the image captioning model\n",
    "attn_model = CaptioningRNN(\n",
    "          cell_type='attention',\n",
    "          word_to_idx=data_dict['vocab']['token_to_idx'],\n",
    "          input_dim=1280, # hard-coded, do not modify\n",
    "          hidden_dim=512,\n",
    "          wordvec_dim=256,\n",
    "          ignore_index=NULL_index,\n",
    "          **to_float_cuda)\n",
    "\n",
    "for learning_rate in [1e-3]:\n",
    "  print('learning rate is: ', learning_rate)\n",
    "  attn_model_submit, attn_loss_submit = captioning_train(attn_model, small_image_data, small_caption_data,\n",
    "                num_epochs=num_epochs, batch_size=batch_size,\n",
    "                learning_rate=learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5ham_O1TG_z7"
   },
   "source": [
    "### Test-time sampling and visualization\n",
    "As with RNN and LSTM, the samples on training data should be very good; the samples on validation data will probably make less sense.\n",
    "\n",
    "We use the `attention_visualizer` function from `eecs598/utils.py` to visualize the attended regions per generated word. Note that sometimes the attended regions (brighter) might not make much sense particially due to our low resolution image input. In real applications, the attended regions are more accurate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0i8KNWSDSLNu"
   },
   "outputs": [],
   "source": [
    "# Sample a minibatch and show the reshaped 112x112 images,\n",
    "# GT captions, and generated captions by your model.\n",
    "\n",
    "batch_size = 3\n",
    "from torchvision.utils import make_grid\n",
    "from torchvision import transforms\n",
    "\n",
    "for split in ['train', 'val']:\n",
    "  sample_idx = torch.randint(0, num_train if split=='train' else num_val, (batch_size,))\n",
    "  sample_images = data_dict[split+'_images'][sample_idx]\n",
    "  sample_captions = data_dict[split+'_captions'][sample_idx]\n",
    "\n",
    "  # decode_captions is loaded from a4_helper.py\n",
    "  gt_captions = decode_captions(sample_captions, data_dict['vocab']['idx_to_token'])\n",
    "  attn_model.eval()\n",
    "  generated_captions, attn_weights_all = attn_model.sample(sample_images)\n",
    "  generated_captions = decode_captions(generated_captions, data_dict['vocab']['idx_to_token'])\n",
    "\n",
    "  for i in range(batch_size):\n",
    "    plt.imshow(sample_images[i].permute(1, 2, 0))\n",
    "    plt.axis('off')\n",
    "    plt.title('%s\\nAttention LSTM Generated:%s\\nGT:%s' % (split, generated_captions[i], gt_captions[i]))\n",
    "    plt.show()\n",
    "    \n",
    "    tokens = generated_captions[i].split(' ')\n",
    "    \n",
    "    vis_attn = []\n",
    "    for j in range(len(tokens)):\n",
    "      img = sample_images[i]\n",
    "      attn_weights = attn_weights_all[i][j]\n",
    "      token = tokens[j]\n",
    "      img_copy = attention_visualizer(img, attn_weights, token)\n",
    "      vis_attn.append(transforms.ToTensor()(img_copy))\n",
    "    \n",
    "    plt.rcParams['figure.figsize'] = (20.0, 20.0)\n",
    "    vis_attn = make_grid(vis_attn, nrow=8)\n",
    "    plt.imshow(torch.flip(vis_attn, dims=(0,)).permute(1, 2, 0))\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "    plt.rcParams['figure.figsize'] = (10.0, 8.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ghB8BwfUpmI5"
   },
   "source": [
    "# Final checks\n",
    "Make sure you run \"Runtime -> Restart and run all...\" to double check the RNN/LSTM code is still functioning well after all the changes you have made!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lgwer0NRQ07Y"
   },
   "source": [
    "# Save results\n",
    "\n",
    "Once all the cells are completed, save the final losses for submission."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-_21vKQeRMpC"
   },
   "outputs": [],
   "source": [
    "submission = {\n",
    "    'rnn_losses': rnn_loss_submit,\n",
    "    'lstm_losses':  lstm_loss_submit,\n",
    "    'attn_losses': attn_loss_submit\n",
    "}\n",
    "submission_path = os.path.join(GOOGLE_DRIVE_PATH, 'rnn_lstm_attention_submission.pkl')\n",
    "dump_results(submission, submission_path)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "rnn_lstm_attention_captioning.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
